package com.xzx.flink.streamapi.state;

import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.CoProcessFunction;
import org.apache.flink.util.Collector;

/**
 * 模拟  SELECT * FROM A INNER JOIN B WHERE A.id = B.id；  这个逻辑
 *
 * @version 1.0
 * @auther xinzhixuan
 * @date 2022/5/14 00:59
 */
public class ListState_01_InnerJoin {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        SingleOutputStreamOperator<Tuple3<String, String, Long>> stream1 = env.fromElements(Tuple3.of("a", "stream-1", 1000L), Tuple3.of("b", "stream-1", 2000L)).assignTimestampsAndWatermarks(WatermarkStrategy.<Tuple3<String, String, Long>>forMonotonousTimestamps().withTimestampAssigner(new SerializableTimestampAssigner<Tuple3<String, String, Long>>() {
            @Override
            public long extractTimestamp(Tuple3<String, String, Long> t, long l) {
                return t.f2;
            }
        }));

        SingleOutputStreamOperator<Tuple3<String, String, Long>> stream2 = env.fromElements(Tuple3.of("a", "stream-2", 3000L), Tuple3.of("b", "stream-2", 4000L)).assignTimestampsAndWatermarks(WatermarkStrategy.<Tuple3<String, String, Long>>forMonotonousTimestamps().withTimestampAssigner(new SerializableTimestampAssigner<Tuple3<String, String, Long>>() {
            @Override
            public long extractTimestamp(Tuple3<String, String, Long> t, long l) {
                return t.f2;
            }
        }));

        stream1.keyBy(d -> d.f0).connect(stream2.keyBy(d -> d.f0)).process(new CoProcessFunction<Tuple3<String, String, Long>, Tuple3<String, String, Long>, String>() {
            private ListState<Tuple3<String, String, Long>> state1;
            private ListState<Tuple3<String, String, Long>> state2;

            @Override
            public void open(Configuration parameters) throws Exception {
                state1 = getRuntimeContext().getListState(new ListStateDescriptor<>("stream1", Types.TUPLE(Types.STRING, Types.STRING, Types.LONG)));
                state2 = getRuntimeContext().getListState(new ListStateDescriptor<>("stream2", Types.TUPLE(Types.STRING, Types.STRING, Types.LONG)));
            }

            @Override
            public void processElement1(Tuple3<String, String, Long> value, CoProcessFunction<Tuple3<String, String, Long>, Tuple3<String, String, Long>, String>.Context ctx, Collector<String> out) throws Exception {
                state1.add(value);
                if (state2.get() != null) {
                    // 这里面左右关联的时候不需要再判断左边跟右边key是否相同了，因为这里这是基于KeyedStream的状态，所以相同的key一定共享一份状态，且不同的key状态隔离的
                    for (Tuple3<String, String, Long> right : state2.get()) {
                        out.collect("left:" + value + "==>right:" + right);
                    }
                }
            }

            @Override
            public void processElement2(Tuple3<String, String, Long> value, CoProcessFunction<Tuple3<String, String, Long>, Tuple3<String, String, Long>, String>.Context ctx, Collector<String> out) throws Exception {
                state2.add(value);
                if (state1.get() != null) {
                    for (Tuple3<String, String, Long> left : state1.get()) {
                        out.collect("left:" + left + "==>right:" + value);
                    }
                }

            }

            @Override
            public void close() throws Exception {
                state1.clear();
                state2.clear();
            }
        }).print();

        env.execute(ListState_01_InnerJoin.class.getSimpleName());
    }
}
